import base64
import copy
import json
import time
import hmac


class JWT:
    def __init__(self):
        pass

    @staticmethod
    def encode(payload, key, exp=300):
        header = {'alg': "HS256", 'typ': 'JWT'}
        # 将字典转化成json字符串
        # dumps 中有个方法将转化完的json字符串去除空格 separators=(",", ":")去除逗号后面以及冒号后面的空格
        # sort_keys=True 将字典排序
        header_json = json.dumps(header, separators=(",", ":"), sort_keys=True)
        # 将字符串转化成字节串
        headers_base64 = JWT.base64_encode(header_json.encode())

        # 初始化payload
        # 将传进来payload+exp两个参数合并  作为payload主体
        # 为了不污染原有字典中的数据需要将字典先拷贝一份
        payload_deepcopy = copy.deepcopy(payload)
        # 将过期时间加入到payload中
        payload_deepcopy["exp"] = exp + time.time()
        payload_json = json.dumps(payload_deepcopy, separators=(",", ":"), sort_keys=True)
        payload_base64 = JWT.base64_encode(payload_json.encode())

        # 签名sign
        # new()中的三个参数key, msg=None, digestmod=''
        hm = hmac.new(key=key.encode(), msg=headers_base64 + b"." + payload_base64, digestmod="SHA256")
        # hm使用二进制转码
        hm_base64 = JWT.base64_encode(hm.digest())
        return headers_base64 + b"." + payload_base64 + b"." + hm_base64

    # 将"="去除
    @staticmethod
    def base64_encode(base64_sign):
        return base64.urlsafe_b64encode(base64_sign).replace(b"=", b"")

    # 将删除的"="补充回来
    @staticmethod
    def base64_decode(base64_token):
        """
        通过base64_token长度除以4如果长度未能除尽说明这里曾经被删除过"="，
        再通过 4-rem 补全"="
        :param base64_token:
        :return:
        """
        rem = len(base64_token) % 4
        if rem > 0:
            base64_token += b"=" * (4 - rem)
        return base64.urlsafe_b64decode(base64_token)

    # 校验数据是否有问题
    @staticmethod
    def decode(token, key):
        # 将token中的header,payload,sign 分离出来
        header_base64, payload_base64, sign_base64 = token.split(b".")
        # 通过分离出来的header_base64,payload_base64,key进行再次hmac
        hm = hmac.new(key.encode(), header_base64 + b"." + payload_base64, digestmod="SHA256")
        if sign_base64 != JWT.base64_encode(hm.digest()):
            raise
        # 检查token是否过期
        payload_json = JWT.base64_decode(payload_base64)
        payload = json.loads(payload_json)
        exp = payload.get("exp")
        if exp - time.time() <= 0:
            print("Token已过期")
            raise
        return payload


if __name__ == '__main__':
    jwt = JWT()
    # 加密
    Token = jwt.encode({"xiaoxiao": '111'},
                       key="eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJleHAiOjE2Mzk2MTgwMTcuMDU0NDU2LCJ4aWFveGlhbyI6IjExMSJ9.3cYgVc4PgPBHwluleL0j5UpWvyk4uqDDNelTJ4pPP1o",
                       exp=300)
    # 校验
    print(jwt.decode(
        token=Token,
        key="eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJleHAiOjE2Mzk2MTgwMTcuMDU0NDU2LCJ4aWFveGlhbyI6IjExMSJ9.3cYgVc4PgPBHwluleL0j5UpWvyk4uqDDNelTJ4pPP1o"))
